import os
import tarfile
import urllib
import zipfile
from pathlib import Path
import numpy as np
from taming.data.helper_types import Annotation
string_classes = str
from tqdm import tqdm


def unpack(path):
    if path.endswith('tar.gz'):
        with tarfile.open(path, 'r:gz') as tar:
            tar.extractall(path=os.path.split(path)[0])
    elif path.endswith('tar'):
        with tarfile.open(path, 'r:') as tar:
            tar.extractall(path=os.path.split(path)[0])
    elif path.endswith('zip'):
        with zipfile.ZipFile(path, 'r') as f:
            f.extractall(path=os.path.split(path)[0])
    else:
        raise NotImplementedError('Unknown file extension: {}'.format(os.
            path.splitext(path)[1]))


def reporthook(bar):
    """tqdm progress bar for downloads."""

    def hook(b=1, bsize=1, tsize=None):
        if tsize is not None:
            bar.total = tsize
        bar.update(b * bsize - bar.n)
    return hook


def get_root(name):
    base = 'data/'
    root = os.path.join(base, name)
    os.makedirs(root, exist_ok=True)
    return root


def is_prepared(root):
    return Path(root).joinpath('.ready').exists()


def mark_prepared(root):
    Path(root).joinpath('.ready').touch()


def prompt_download(file_, source, target_dir, content_dir=None):
    targetpath = os.path.join(target_dir, file_)
    while not os.path.exists(targetpath):
        if content_dir is not None and os.path.exists(os.path.join(
            target_dir, content_dir)):
            break
        print("Please download '{}' from '{}' to '{}'.".format(file_,
            source, targetpath))
        if content_dir is not None:
            print("Or place its content into '{}'.".format(os.path.join(
                target_dir, content_dir)))
        input('Press Enter when done...')
    return targetpath


def download_url(file_, url, target_dir):
    targetpath = os.path.join(target_dir, file_)
    os.makedirs(target_dir, exist_ok=True)
    with tqdm(unit='B', unit_scale=True, unit_divisor=1024, miniters=1,
        desc=file_) as bar:
        urllib.request.urlretrieve(url, targetpath, reporthook=reporthook(bar))
    return targetpath


def download_urls(urls, target_dir):
    paths = dict()
    for fname, url in urls.items():
        outpath = download_url(fname, url, target_dir)
        paths[fname] = outpath
    return paths


def quadratic_crop(x, bbox, alpha=1.0):
    """bbox is xmin, ymin, xmax, ymax"""
    im_h, im_w = tuple(x.shape)[:2]
    bbox = np.array(bbox, dtype=np.float32)
    bbox = np.clip(bbox, 0, max(im_h, im_w))
    center = 0.5 * (bbox[0] + bbox[2]), 0.5 * (bbox[1] + bbox[3])
    w = bbox[2] - bbox[0]
    h = bbox[3] - bbox[1]
    l = int(alpha * max(w, h))
    l = max(l, 2)
    required_padding = -1 * min(center[0] - l, center[1] - l, im_w - (
        center[0] + l), im_h - (center[1] + l))
    required_padding = int(np.ceil(required_padding))
    if required_padding > 0:
        padding = [[required_padding, required_padding], [required_padding,
            required_padding]]
        padding += [[0, 0]] * (len(tuple(x.shape)) - 2)
        x = np.pad(x, padding, 'reflect')
        center = center[0] + required_padding, center[1] + required_padding
    xmin = int(center[0] - l / 2)
    ymin = int(center[1] - l / 2)
    return np.array(x[ymin:ymin + l, xmin:xmin + l, ...])




